from pysnmp.hlapi import *
from datetime import datetime
import logging

logger = logging.getLogger(__name__)

class SNMPClient:
    def __init__(self, host, community='public', port=161, version='2c'):
        self.host = host
        self.community = community
        self.port = port
        self.version = version
        logger.info(f"初始化SNMP客户端: host={host}, community={community}, port={port}, version={version}")

    def get_snmp_data(self, oid):
        """获取单个 OID 的数据"""
        try:
            logger.debug(f"尝试获取OID {oid} 的数据")
            
            # 创建SNMP引擎
            snmp_engine = SnmpEngine()
            # 设置认证信息
            auth_data = CommunityData(self.community, mpModel=1)
            # 设置目标主机
            target = UdpTransportTarget((self.host, self.port), timeout=2, retries=2)
            # 设置上下文
            context = ContextData()
            # 设置OID
            object_identity = ObjectIdentity(oid)

            # 执行SNMP GET请求
            error_indication, error_status, error_index, var_binds = next(
                getCmd(snmp_engine, auth_data, target, context,
                      ObjectType(object_identity))
            )

            if error_indication:
                logger.error(f"SNMP错误: {error_indication}")
                raise Exception(f'SNMP错误: {error_indication}')
            elif error_status:
                logger.error(f"SNMP错误状态: {error_status}")
                raise Exception(f'SNMP错误: {error_status}')
            
            # 获取返回值
            for var_bind in var_binds:
                value = var_bind[1]
                logger.debug(f"获取到OID {oid} 的值: {value}")
                return str(value)
            return None

        except Exception as e:
            logger.error(f"获取SNMP数据失败: {str(e)}")
            raise Exception(f'SNMP获取数据失败: {str(e)}')

    def get_disk_usage(self):
        """获取磁盘使用情况"""
        try:
            # 使用 UCD-SNMP-MIB 的磁盘OID
            DISK_PATH_OID = '1.3.6.1.4.1.2021.9.1.2'    # 磁盘路径
            DISK_TOTAL_OID = '1.3.6.1.4.1.2021.9.1.5'   # 总大小 (kB)
            DISK_USED_OID = '1.3.6.1.4.1.2021.9.1.6'    # 已使用 (kB)
            DISK_PERCENT_OID = '1.3.6.1.4.1.2021.9.1.9'  # 使用百分比

            # 创建SNMP引擎和其他参数
            snmp_engine = SnmpEngine()
            auth_data = CommunityData(self.community, mpModel=1)
            target = UdpTransportTarget((self.host, self.port))
            context = ContextData()

            # 获取所有磁盘分区
            storage_data = []
            
            # 先尝试直接获取使用百分比
            for (error_indication, error_status, error_index, var_binds) in nextCmd(
                snmp_engine,
                auth_data,
                target,
                context,
                ObjectType(ObjectIdentity(DISK_PATH_OID)),
                ObjectType(ObjectIdentity(DISK_PERCENT_OID)),
                lexicographicMode=False
            ):
                if error_indication or error_status:
                    continue

                path = str(var_binds[0][1])
                usage = int(var_binds[1][1])
                
                if path == "/" or path.startswith("/dev"):  # 只关注根分区和实际设备
                    storage_data.append({
                        'path': path,
                        'usage': usage
                    })
                    logger.info(f"发现磁盘分区: {path}, 使用率: {usage}%")

            # 如果找到磁盘，返回最大使用率
            if storage_data:
                max_usage = max(item['usage'] for item in storage_data)
                logger.info(f"磁盘使用率: {max_usage}% (最大值)")
                return round(max_usage, 2)

            # 如果上面的方法失败，尝试使用总量和已用量计算
            logger.info("尝试使用总量和已用量计算磁盘使用率")
            for (error_indication, error_status, error_index, var_binds) in nextCmd(
                snmp_engine,
                auth_data,
                target,
                context,
                ObjectType(ObjectIdentity(DISK_PATH_OID)),
                ObjectType(ObjectIdentity(DISK_TOTAL_OID)),
                ObjectType(ObjectIdentity(DISK_USED_OID)),
                lexicographicMode=False
            ):
                if error_indication or error_status:
                    continue

                path = str(var_binds[0][1])
                total = int(var_binds[1][1])
                used = int(var_binds[2][1])

                if (path == "/" or path.startswith("/dev")) and total > 0:
                    usage = (used / total) * 100
                    storage_data.append({
                        'path': path,
                        'usage': usage
                    })
                    logger.info(f"发现磁盘分区: {path}, 使用率: {usage}%")

            if storage_data:
                max_usage = max(item['usage'] for item in storage_data)
                logger.info(f"磁盘使用率: {max_usage}% (最大值)")
                return round(max_usage, 2)
            
            logger.warning("未找到有效的磁盘存储")
            return 0

        except Exception as e:
            logger.error(f"获取磁盘使用率失败: {str(e)}")
            return 0

    def get_system_metrics(self):
        """获取系统指标"""
        try:
            logger.info(f"开始获取系统指标: {self.host}")
            
            # Linux系统的标准OID
            CPU_LOAD_OID = '1.3.6.1.4.1.2021.11.9.0'  # CPU使用率
            TOTAL_MEMORY_OID = '1.3.6.1.4.1.2021.4.5.0'  # 总内存
            FREE_MEMORY_OID = '1.3.6.1.4.1.2021.4.6.0'   # 空闲内存
            BUFFER_MEMORY_OID = '1.3.6.1.4.1.2021.4.14.0'  # 缓冲区内存
            CACHED_MEMORY_OID = '1.3.6.1.4.1.2021.4.15.0'  # 缓存内存
            DISK_TOTAL_OID = '1.3.6.1.4.1.2021.9.1.6.1'  # 总磁盘空间
            DISK_USED_OID = '1.3.6.1.4.1.2021.9.1.8.1'   # 已用磁盘空间

            # 获取CPU使用率
            try:
                cpu_raw = float(self.get_snmp_data(CPU_LOAD_OID) or 0)
                cpu_usage = min(cpu_raw, 100)  # 确保不超过100%
                logger.info(f"CPU使用率: {cpu_usage}%")
            except Exception as e:
                logger.error(f"获取CPU使用率失败: {str(e)}")
                cpu_usage = 0

            # 获取内存使用情况
            try:
                total_memory = int(self.get_snmp_data(TOTAL_MEMORY_OID) or 0)
                free_memory = int(self.get_snmp_data(FREE_MEMORY_OID) or 0)
                buffer_memory = int(self.get_snmp_data(BUFFER_MEMORY_OID) or 0)
                cached_memory = int(self.get_snmp_data(CACHED_MEMORY_OID) or 0)
                
                if total_memory > 0:
                    used_memory = total_memory - free_memory - buffer_memory - cached_memory
                    memory_usage = (used_memory / total_memory) * 100
                    logger.info(f"内存使用率: {memory_usage}%")
                else:
                    memory_usage = 0
                    logger.warning("无法获取内存信息")
            except Exception as e:
                logger.error(f"获取内存使用率失败: {str(e)}")
                memory_usage = 0

            # 获取磁盘使用情况
            disk_usage = self.get_disk_usage()

            metrics = {
                'cpu_usage': round(cpu_usage, 2),
                'memory_usage': round(memory_usage, 2),
                'disk_usage': disk_usage,
                'timestamp': datetime.now()
            }
            
            logger.info(f"系统指标获取完成: {metrics}")
            return metrics
            
        except Exception as e:
            logger.error(f"获取系统指标失败: {str(e)}")
            raise Exception(f'获取系统指标失败: {str(e)}')

    def test_connection(self):
        """测试SNMP连接"""
        try:
            logger.info(f"测试SNMP连接: {self.host}")
            SYSTEM_DESCRIPTION_OID = '1.3.6.1.2.1.1.1.0'
            result = self.get_snmp_data(SYSTEM_DESCRIPTION_OID)
            success = result is not None
            logger.info(f"SNMP连接测试{'成功' if success else '失败'}: {result}")
            return success
        except Exception as e:
            logger.error(f"SNMP连接测试失败: {str(e)}")
            return False

    def get_windows_metrics(self):
        """获取Windows系统指标"""
        try:
            # Windows系统的OID
            CPU_LOAD_OID = '1.3.6.1.2.1.25.3.3.1.2.1'  # CPU使用率
            MEMORY_TOTAL_OID = '1.3.6.1.2.1.25.2.2.0'  # 总内存
            MEMORY_USED_OID = '1.3.6.1.2.1.25.2.3.1.6.1'  # 已用内存
            DISK_TOTAL_OID = '1.3.6.1.2.1.25.2.3.1.5.1'  # 总磁盘空间
            DISK_USED_OID = '1.3.6.1.2.1.25.2.3.1.6.1'  # 已用磁盘空间

            # 获取CPU使用率
            try:
                cpu_raw = int(self.get_snmp_data(CPU_LOAD_OID) or 0)
                cpu_usage = min(cpu_raw, 100)
            except:
                cpu_usage = 0

            # 获取内存使用情况
            try:
                total_memory = int(self.get_snmp_data(MEMORY_TOTAL_OID) or 0)
                used_memory = int(self.get_snmp_data(MEMORY_USED_OID) or 0)
                memory_usage = (used_memory / total_memory * 100) if total_memory > 0 else 0
            except:
                memory_usage = 0

            # 获取磁盘使用情况
            try:
                disk_total = int(self.get_snmp_data(DISK_TOTAL_OID) or 0)
                disk_used = int(self.get_snmp_data(DISK_USED_OID) or 0)
                disk_usage = (disk_used / disk_total * 100) if disk_total > 0 else 0
            except:
                disk_usage = 0

            return {
                'cpu_usage': round(cpu_usage, 2),
                'memory_usage': round(memory_usage, 2),
                'disk_usage': round(disk_usage, 2),
                'timestamp': datetime.now()
            }
        except Exception as e:
            raise Exception(f'获取Windows系统指标失败: {str(e)}')

    def detect_system_type(self):
        """检测系统类型"""
        try:
            SYSTEM_DESCRIPTION_OID = '1.3.6.1.2.1.1.1.0'
            description = str(self.get_snmp_data(SYSTEM_DESCRIPTION_OID)).lower()
            
            if 'windows' in description:
                return 'windows'
            elif 'linux' in description:
                return 'linux'
            else:
                return 'unknown'
        except:
            return 'unknown'

    def get_metrics(self):
        """根据系统类型获取相应的指标"""
        system_type = self.detect_system_type()
        
        if system_type == 'windows':
            return self.get_windows_metrics()
        else:
            return self.get_system_metrics()  # 默认使用Linux/Unix的OID 